import unittest

import requests


class TestKnowledgeBaseAPI(unittest.TestCase):
    BASE_URL = "http://180.184.65.98:38880/atomgit"

    def test_search_papers(self):
        """
        测试根据文本查询搜索论文的API。
        """
        url = f"{self.BASE_URL}/search_papers"
        params = {"query": "机器学习", "top_k": 5}
        response = requests.get(url, params=params)
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIsInstance(data, list)
        self.assertLessEqual(len(data), 5)
        # 根据实际响应结构添加更多断言

    def test_query_by_paper_id(self):
        """
        测试根据论文ID查询论文的API。
        """
        url = f"{self.BASE_URL}/query_by_paper_id"
        params = {"paper_id": "655d674d939a5f4082e02457", "top_k": 5}
        response = requests.get(url, params=params)
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIsInstance(data, list)
        self.assertLessEqual(len(data), 5)
        # 根据实际响应结构添加更多断言

    def test_query_by_title(self):
        """
        测试根据论文标题查询论文的API。
        """
        url = f"{self.BASE_URL}/query_by_title"
        params = {"title": "'MultiDelete for Multimodal Machine Unlearning'", "top_k": 5}
        response = requests.get(url, params=params)
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIsInstance(data, list)
        self.assertLessEqual(len(data), 5)
        # 根据实际响应结构添加更多断言

    def test_metadata(self):
        """
        测试获取论文数据库元数据信息的API。
        """
        url = f"{self.BASE_URL}/metadata"
        response = requests.get(url)
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIsInstance(data, dict)

    def test_query_by_title_contain(self):
        """
        测试搜索标题中包含特定文本的论文片段的API。
        """
        url = f"{self.BASE_URL}/query_by_title_contain"
        params = {"title": "机器学习", "top_k": 5}
        response = requests.get(url, params=params)
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIsInstance(data, list)
        self.assertLessEqual(len(data), 5)

    def test_query_by_chunk_contain(self):
        """
        测试搜索论文内容片段中包含特定文本的API。
        """
        url = f"{self.BASE_URL}/query_by_chunk_contain"
        params = {"chunk": "深度学习", "top_k": 5}
        response = requests.get(url, params=params)
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIsInstance(data, list)
        self.assertLessEqual(len(data), 5)


if __name__ == "__main__":
    unittest.main()
